-------------------------------------------------------------------------------
-- Async IO - A module to provide a 'blocking' interface for non-blocking IO
-- 
-- This allows programs written with blocking assumption to continue to
-- run when the underlying becomes non-blocking thus can be scheduled
-- with a dispatcher like select() or libevent 
-- 
--
-- asyncio.wrap(io_object, buffered_write)
--
-- returns an IO object that can be used in place of normal file or socket
-- but would yield to the dispatcher on timeout and continue when resumed
--
-- In other words, the returned object must be used in coroutines(or functions
-- called in them). There is no point using this if there is no coroutine
-- dispatcher.
--
-- Copyright 2007 - Gary Ng<linux@garyng.com>
-- This code can be distributed under LGPL license
--
-------------------------------------------------------------------------------
local type=type
local rawget = rawget
local coyield,costatus = coroutine.yield,coroutine.status
local coresume,corunning = coroutine.resume,coroutine.running
local tjoin = table.concat
local time = os.time
local duration = os.difftime
local tostring = tostring

module('asyncio', package.seeall)
STALLED_SPIN = 20
CHECK_STALLED = true

local function proto_index(o, k)  

    local v = rawget(o, '_mt')[k]
    if v then return v end

    local base = rawget(o, '_proto')
    if base then
      v = base[k]
      if type(v) == "function" then 
          return function(x,...) return v(base,...) end  
      else
          return v
      end
    end
end

local asyncio={}; asyncio.__index=asyncio
local asyncio_meta = {__index = proto_index,  __mode="v",__tostring=function(o) return 'asyncio:'..tostring(o._proto) end }

function wrap(conn, buffered_write)
    if buffered_write == nil then
        if conn.buffered_write == nil then buffered_write = 1
        else buffered_write = conn.buffered_write end
    end
    if conn.settimeout then conn:settimeout(0) end
    return setmetatable({_proto=conn, _mt = asyncio, active = 1, buffered_write = buffered_write}, asyncio_meta) 
end

local function readline(conn, read_fn, part, timeout)
   local d,e,c,q
   local start = time()  
   local skt,abort
   local suspended=0
   local dangling
   part = part or ""
   timeout = timeout or 0
   repeat
        local nl = part:find('\n')
        if not nl then
            d,e,c,q = read_fn(conn, 128)
            part = part .. (d or c or "")
            if not part:find('\n') then
                if e == "timeout" then
                    if not conn.love_live and dangling and suspended > STALLED_SPIN then return nil, "stalled", "" end
                    if CHECK_STALLED and suspended > 15000 then return nil, "stalled","" end 
                    if timeout==0 or duration(time(), start) < timeout then 
                      suspended = suspended + 1
                      if q == "write" then
                        skt,abort = coyield(conn, "writing") 
                        if abort then dangling = true end
                      else
                        skt,abort = coyield(conn, "reading") 
                        if abort then dangling = true end
                      end
                    else 
                      return nil, e, part 
                    end
                elseif e then return nil, e, part 
                else suspended = 0
                end
            end
        else
            return part:sub(1, nl - 1), nil, part:sub(nl + 1)
        end
   until false
end

local function receivefrom(conn, read_fn, size, timeout)
   local start = time()   
   local skt,abort
   timeout = timeout or 0
   repeat
        local a,b,c,q = read_fn(conn, size)
        if a then return a,b,c
        elseif b == "timeout" then 
            if timeout==0 or duration(time(), start) < timeout then 
              if q == "write" then
                skt,abort = coyield(conn, "writing") 
                if abort then return nil,"cancelled","" end
              else
                skt,abort = coyield(conn, "reading") 
                if abort then return nil,"cancelled","" end
              end
            else return nil, b, "" end
        elseif b then return nil, b, "" end
   until false
end

local function udpreceive(conn, read_fn, size, timeout)
   local start = time()   
   local skt,abort
   timeout=timeout or 0
   repeat
        local a,b,c,q = read_fn(conn, size)
        if a then return a
        elseif b == "timeout" then 
            if timeout==0 or duration(time(), start) < timeout then 
              if q == "write" then
                skt,abort = coyield(conn, "writing") 
                if abort then return nil,"cancelled","" end
              else
                skt,abort = coyield(conn, "reading") 
                if abort then return nil,"cancelled","" end
              end
            else return nil, b,"" end
        elseif b then return nil, b,"" end
   until false
end

local function read(conn, read_fn, size, part, timeout)
   local skt,abort
   part = part or ""
   local t={part}
   size = size - #part
   local start = time()   
   local suspended=0
   local dangling
   timeout = timeout or 0
   repeat
        if math.random(100) > 95 then
          skt,abort = coyield(conn, "reading")
        end
        local d,e,c,q = read_fn(conn, size)
        local data = d or c or ""
        size = size - #data
        t[#t+1] = data
        if size == 0 then return tjoin(t,"") 
        elseif e == "timeout" then 
            if not conn.love_live and dangling and suspended > STALLED_SPIN then return nil, "stalled", "" end
            if CHECK_STALLED and suspended > 15000 then return nil, "stalled","" end 
            --assert(suspended < 10000)
            if timeout==0 or duration(time(), start) < timeout then 
              suspended = suspended + 1
              if q == "write" then
                skt,abort = coyield(conn, "writing") 
                if abort then dangling = true end
              else
                skt,abort = coyield(conn, "reading") 
                if abort then dangling = true end
              end
            else 
              local part = tjoin(t,"")
              return nil, e, part
            end
        elseif e then 
          return nil, e, tjoin(t,"") 
        else
          suspended = 0
        end
   until false
end

local function readall(conn, read_fn, part, timeout)
    local t={}
    part = part or ""
    repeat
        local d,e
        d,e,part = read(conn, read_fn, 8192, part, timeout)
        local data = d or part or ""
        t[#t+1] = data
        if not e then return tjoin(t,"") 
        else 
          local part = tjoin(t,"")
          return nil,e,part
        end
    until false
end

local function write(conn, write_fn, data, buffered_write, timeout)
   local c,e,i,j,q
   i = 0
   local chunk = data
   local start = time()  
   local written = 0
   local skt,abort 
   local dangling
   local suspended = 0
   timeout = timeout or 0
   repeat
        if (math.random(100) > 95) then
          skt,abort = coyield(conn, "writing")
          if abort then return nil,"cancelled","" end
        end
        if buffered_write then
            c,e,i,q = write_fn(conn, chunk, i+1)
            written = c or i
        else
            i,e,q = write_fn(conn, chunk:sub(i+1))
            written = written + i
        end
        if e == "timeout" then 
            if not conn.love_live and dangling and suspended > STALLED_SPIN then return nil, "stalled", "" end
            if CHECK_STALLED and suspended > 15000 then return nil, "stalled","" end 
            --assert(suspended < 10000)
            if timeout==0 or duration(time(), start) < timeout then 
              suspended = suspended + 1
              if q == "read" then
                skt,abort = coyield(conn, "reading") 
                if abort then dangling = true end
              else
                skt,abort = coyield(conn, "writing") 
                if abort then dangling = true end
              end
            else 
              return nil, e, written  
            end
        elseif e then return c,e,written 
        else return written end
   until false
end


function asyncio:_rawreceivefrom(size)
    local read_fn = self._proto.receivefrom
    return read_fn(self._proto, size)
end

function asyncio:_rawread(size)
    local read_fn = self._proto.read or self._proto.receive
    return read_fn(self._proto, size)
end

function asyncio:_rawwrite(data, i, j)
    local write_fn = self._proto.write or self._proto.send
    return write_fn(self._proto, data, i, j)
end

function asyncio:receivefrom(size)
    local skt,abort
    local read_fn = self._rawreceivefrom
    --serialize in thread unit
    repeat
        local p = self.pending_read
        if p and costatus(p) == "suspended" then 
            skt,abort = coyield(self, "reading") 
            if abort then return nil,"cancelled","" end
        else
            break
        end
    until false

    self.pending_reading = corunning()
    return receivefrom(self, read_fn, size, self.timeout)
end

function asyncio:read(pattern, partial)
    local skt,abort
    if partial and type(partial) == "string" then self._part = (self._part or "").. partial end
    local isudp = self._proto.sendto
    local read_fn = self._rawread
    if not isudp then
      pattern = pattern or "*l"
    end
    --serialize in thread unit
    repeat
        local p = self.pending_read
        if p and costatus(p) == "suspended" then 
            skt,abort = coyield(self, "reading") 
            if abort then return nil,"cancelled","" end
        else
            break
        end
    until false

    self.pending_reading = corunning()
    if type(pattern) == "number" then 
        if isudp then 
          return udpreceive(self, read_fn, pattern, self.timeout) 
        else
          local x, e, part = read(self, read_fn, pattern, self._part, self.timeout)
          self._part = part
          self.pending_reading = nil
          return x, e, part
        end
    elseif pattern == "*l" then 
        local x, e, part = readline(self, read_fn, self._part, self.timeout)
        self._part = part
        self.pending_reading = nil
        if x then x = x:match("^(.-)\r*$") end
        return x, e, part
    elseif pattern == "*a" or pattern == "*all" then 
        local x, e, part =  readall(self,read_fn,self._part, self.timeout)
        self._part = nil
        self.pending_reading = nil
        return x or part or "", e, ""
    else 
        self.pending_reading = nil
        error("unknown pattern for read")
    end

end
asyncio.receive = asyncio.read

function asyncio:connect(...)
  local skt,abort
  local start = time()
  self._proto:settimeout(self.timeout or 0.5)
  repeat
	  local ret,err = self._proto:connect(...)
    if err == "timeout" and (not self.timeout or duration(time(), start) < self.timeout) then 
      skt,abort = coyield(self, "writing") 
      if abort then return nil,"canceled","" end
    else
      return ret, err,""
    end
  until false
end

function asyncio:dirty()
    local part = self._part
    return part and #part > 0 or self._proto:dirty()
end

function asyncio:close()
    self.active = self.active - 1
    if self.active == 0 then
        local proto = self._proto
        --self._proto = nil
        return proto:close()
    end
end
function asyncio:open()
    self.active = self.active + 1
    return self
end

function asyncio:write(data)
    local skt,abort 
    local write_fn = self._rawwrite
    
    --serialize in thread unit
    repeat
        local p = self.pending_write
        if p and costatus(p) == "suspended" then 
            skt,abort = coyield(self, "writing") 
            if abort then return nil,"cancelled","" end
        else
            break
        end
    until false

    self.pending_write = corunning()
    local a,b,c =  write(self, write_fn, data, self.buffered_write, self.timeout)
    self.pending_write = nil
    return a,b,c

end
asyncio.send = asyncio.write

function asyncio:settimeout(timeout)
    self.timeout = timeout
end
